# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import atexit
import asyncio
import enum
import json
import logging
import os
import signal
import subprocess
import sys

import psutil
from tornado.ioloop import IOLoop
from tornado.process import Subprocess as TornadoSubprocess

from coin.util.slurm.base import TornadoSubprocessWatcher


class SlurmDriverError(Exception):
  pass


def _kill_salloc(salloc_proc):
  salloc_proc.proc.kill()
  salloc_proc.proc.wait()


def _stop_salloc(salloc_proc, wait=True):
  try:
    ps_process = psutil.Process(salloc_proc.proc.pid)
  except psutil.NoSuchProcess:
    return

  for child in ps_process.children(recursive=True):
    os.kill(child.pid, signal.SIGINT)

  if wait:
    salloc_proc.proc.wait()


async def _launch_salloc(argv=None):
  SALLOC_ENTRYPOINT = 'coin/util/slurm/_salloc_entry.py'

  argv = argv or []
  if '-I' not in argv:
    argv += ['-I']

  salloc_proc = TornadoSubprocess(['salloc', *argv, SALLOC_ENTRYPOINT],
                                  stderr=TornadoSubprocess.STREAM,
                                  stdout=TornadoSubprocess.STREAM)
  try:
    line = await salloc_proc.stderr.read_until(b'\n')
    line = line.decode('utf-8').strip()
    if not line.startswith('salloc: Granted job allocation'):
      raise SlurmDriverError(line)
    return salloc_proc
  except Exception:
    _kill_salloc(salloc_proc)
    raise


async def _wait_salloc_env(salloc_proc, check_essential_keys=True):
  try:
    slurm_env_str = await salloc_proc.stdout.read_until(b'\n}')
    slurm_env_str = slurm_env_str.decode('utf-8').strip()
    slurm_env = json.loads(slurm_env_str)
    if check_essential_keys:
      assert 'SLURM_JOB_ID' in slurm_env
    return slurm_env
  except Exception:
    _stop_salloc(salloc_proc)
    raise


def cancel_slurm_job(job_id):
  subprocess.call(['scancel', str(job_id)])


class SlurmDriver(object):
  class State(enum.Enum):
    INIT = 0
    BOOTSTRAP = 1
    READY = 2
    STOP_REQUESTED = 3
    FINISHED = 4
    ERROR = -1

  @staticmethod
  def initialize():
    TornadoSubprocessWatcher.install()

  def __init__(self):
    self._state = SlurmDriver.State.INIT
    self._ioloop = None
    self._salloc_proc = None
    self._slurm_env = None

    self._error_callback = None
    self._ready_future = None
    self._stop_future = None

  @property
  def state(self):
    return self._state

  @property
  def slurm_env(self):
    return self._slurm_env

  @property
  def slurm_job_id(self):
    return self._slurm_env['SLURM_JOB_ID']

  def start(self, salloc_argv=None, ioloop=None, error_callback=None):
    """Run salloc and return an asyncio.Future which resolves
    when the salloc is up and running.
    """
    assert self._state == SlurmDriver.State.INIT

    self._ioloop = ioloop or IOLoop.current()
    self._state = SlurmDriver.State.BOOTSTRAP
    self._error_callback = error_callback
    self._ready_future = asyncio.Future()
    self._ioloop.add_callback(self._salloc_main, salloc_argv)
    return self._ready_future

  def stop(self):
    if self._state in (SlurmDriver.State.FINISHED, SlurmDriver.State.ERROR):
      return self._stop_future

    assert self._state == SlurmDriver.State.READY
    # scancel
    # wait for completion and stop
    # self._ioloop.stop()

    self._state = SlurmDriver.State.STOP_REQUESTED
    assert self._stop_future
    _stop_salloc(self._salloc_proc, wait=False)
    return self._stop_future

  def srun(self, srun_argv, env=None):
    if self._state == SlurmDriver.State.ERROR:
      raise SlurmDriverError('Job terminated')
    assert self._state == SlurmDriver.State.READY

    env = env or dict(os.environ)
    env.update(self._slurm_env)

    future = asyncio.Future()
    self._ioloop.add_callback(self._srun_main, srun_argv=srun_argv, env=env, future=future)
    return future

  async def _srun_main(self, srun_argv, env, future):
    try:
      srun_proc = TornadoSubprocess(['srun', *srun_argv], env=env)
      await srun_proc.wait_for_exit()
      future.set_result(None)
    except Exception as e:
      future.set_exception(e)

  async def _salloc_main(self, salloc_argv):
    assert self._state == SlurmDriver.State.BOOTSTRAP
    try:
      self._salloc_proc = await _launch_salloc(argv=salloc_argv)
      self._slurm_env = await _wait_salloc_env(self._salloc_proc)
    except Exception as e:
      # On exception, resource will be relinquished by
      # _launch_salloc() or _wait_salloc_env().
      self._state = SlurmDriver.State.ERROR
      self._ready_future.set_exception(e)
      return

    try:
      atexit.register(cancel_slurm_job, self._slurm_env['SLURM_JOB_ID'])

      self._state = SlurmDriver.State.READY
      self._ready_future.set_result(None)
      self._stop_future = asyncio.Future()

      self._dump_slurm_env()

      await self._salloc_proc.wait_for_exit(raise_error=False)
      if self._state != SlurmDriver.State.STOP_REQUESTED:
        raise SlurmDriverError('sallow unintentionally exited.')

      self._state = SlurmDriver.State.FINISHED
      if self._salloc_proc.returncode != 0:
        logging.warning('Return code is not zero (%d)', self._salloc_proc.returncode)

    except Exception as e:
      self._state = SlurmDriver.State.ERROR
      if self._error_callback:
        self._ioloop.add_callback(self._error_callback, *sys.exc_info())
      raise

    finally:
      cancel_slurm_job(self._slurm_env['SLURM_JOB_ID'])
      self._stop_future.set_result(None)

  def _dump_slurm_env(self):
    lines = []
    for key in sorted(self._slurm_env.keys()):
      lines.append('%s: %s' % (key, self._slurm_env[key]))
    logging.debug('Slurm environment variables:\n%s', '\n'.join(lines))


'''
Due to its nature, test code for SlurmDriver is hard to be implemented.
So the author leaves a simple test code here.
---

async def test():
  driver = SlurmDriver()
  await driver.start(salloc_argv=['-n', '30', '--spread-job'])
  print('READY!')

  futs = []
  for i in range(0, 100):
    f = driver.srun(['-N1', '-n1', '--exclusive', 'domainname', '-f'])
    futs.append(f)
    print('srun', i)

  for idx, f in enumerate(futs):
    try:
      await f
    except subprocess.CalledProcessError:
      print(idx, 'Failed')

  await driver.stop()

IOLoop.current().run_sync(test)
'''
